use std::{collections::HashMap, sync::Arc};

use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{InitializeRequestParam, InitializeResult, ServerInfo};
use rmcp::service::RequestContext;
use rmcp::{Error as McpError, RoleServer, model::*};
use rmcp::{
    ServerHandler,
    model::{CallToolRequestParam, CallToolResult, Tool as McpTool},
    service::ServerSink,
};
use serde_json::Value;

#[async_trait]
pub trait Tool: Send + Sync {
    fn name(&self) -> String;
    fn description(&self) -> String;
    fn parameters(&self) -> Value;
    async fn call(&self, args: Value) -> Result<CallToolResult>;
    fn to_tool(&self) -> McpTool;
}

pub struct McpToolAdapter {
    tool: McpTool,
    server: ServerSink,
}

impl McpToolAdapter {
    pub fn new(tool: McpTool, server: ServerSink) -> Self {
        Self { tool, server }
    }
}

#[async_trait]
impl Tool for McpToolAdapter {
    fn name(&self) -> String {
        self.tool.name.clone().to_string()
    }

    fn description(&self) -> String {
        self.tool
            .description
            .clone()
            .unwrap_or_default()
            .to_string()
    }

    fn parameters(&self) -> Value {
        serde_json::to_value(&self.tool.input_schema).unwrap_or(serde_json::json!({}))
    }

    async fn call(&self, args: Value) -> Result<CallToolResult> {
        let arguments = match args {
            Value::Object(map) => Some(map),
            _ => None,
        };
        println!("arguments: {:?}", arguments);
        let call_result = self
            .server
            .call_tool(CallToolRequestParam {
                name: self.tool.name.clone(),
                arguments,
            })
            .await?;

        Ok(call_result)
    }

    fn to_tool(&self) -> McpTool {
        self.tool.clone()
    }
}

#[derive(Default)]
pub struct ToolSet {
    tools: HashMap<String, Arc<dyn Tool>>,
}

impl ToolSet {
    pub fn add_tool<T: Tool + 'static>(&mut self, tool: T) {
        self.tools.insert(tool.name(), Arc::new(tool));
    }

    pub fn get_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tools.get(name).cloned()
    }

    pub fn tools(&self) -> Vec<Arc<dyn Tool>> {
        self.tools.values().cloned().collect()
    }

    pub fn clone(&self) -> Self {
        ToolSet {
            tools: self.tools.clone(),
        }
    }

    pub async fn call_tool(&self, tool_name: &str, args: Value) -> Result<CallToolResult> {
        match self.tools.get(tool_name) {
            Some(tool) => {
                let result = tool.call(args).await;
                tracing::debug!("result: {:?}", result);
                return result;
            }
            None => {
                tracing::error!("tool not found: {}", tool_name);
                return Err(anyhow::Error::msg("tool not found"));
            }
        }
    }
}

impl ServerHandler for ToolSet {
    async fn initialize(
        &self,
        _request: InitializeRequestParam,
        context: RequestContext<RoleServer>,
    ) -> Result<InitializeResult, McpError> {
        if let Some(http_request_part) = context.extensions.get::<axum::http::request::Parts>() {
            let initialize_headers = &http_request_part.headers;
            let initialize_uri = &http_request_part.uri;
            tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server");
        }
        Ok(self.get_info())
    }

    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2024_11_05,
            capabilities: ServerCapabilities::builder()
                .enable_prompts()
                .enable_resources()
                .enable_tools()
                .build(),
            server_info: Implementation::from_build_env(),
            instructions: Some("This server provides a counter tool that can increment and decrement values. The counter starts at 0 and can be modified using the 'increment' and 'decrement' tools. Use 'get_value' to check the current count.".to_string()),
        }
    }

    async fn list_tools(
        &self,
        _: Option<::rmcp::model::PaginatedRequestParam>,
        _: ::rmcp::service::RequestContext<::rmcp::service::RoleServer>,
    ) -> Result<::rmcp::model::ListToolsResult, ::rmcp::Error> {
        let mut tools = Vec::new();
        for (_, t) in self.tools.iter() {
            tools.push(t.to_tool());
        }

        Ok(::rmcp::model::ListToolsResult {
            next_cursor: None,
            tools,
        })
    }

    async fn call_tool(
        &self,
        call_tool_request_param: ::rmcp::model::CallToolRequestParam,
        _: ::rmcp::service::RequestContext<::rmcp::service::RoleServer>,
    ) -> Result<::rmcp::model::CallToolResult, rmcp::Error> {
        let res = self
            .call_tool(
                &call_tool_request_param.name,
                call_tool_request_param.arguments.into(),
            )
            .await;
        match res {
            Ok(r) => {
                println!("call tool result: {:?}", r);
                return Ok(r);
            }
            Err(e) => {
                println!("call tool error: {:?}", e);
                return Err(rmcp::Error::invalid_params("hhh", None));
            }
        }
    }
}
